#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/2/5
# @USER    : Shengji He
# @File    : denoising.py
# @Software: PyCharm
# @Version  : Python-
# @TASK: Code for "Blind restoration of a JPEG-compressed image" and "Blind image denoising" figures. Select fname below to switch between the two.
#
# To see overfitting set num_iter to a large value.
# Import libs

from __future__ import print_function
import matplotlib
# matplotlib.use('Agg')
import matplotlib.pyplot as plt
font = {
    'family': 'Times New Roman',
    'style': 'normal',
    'weight': 'bold',
    # 'weight': 'normal',
    # 'color':'black',
    'size': 12
        }

linewidth = 1.5

import seaborn as sns
color_list = sns.color_palette('deep') + sns.color_palette('bright')

import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '3'

import numpy as np
from models import *

import torch
import torch.optim
import torchvision
from skimage.metrics import peak_signal_noise_ratio as compare_psnr
from utils.denoising_utils import *

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True


class Configuration:
    device = torch.device("cpu")
    pin_memory = False
    enable_amp = False
    train_on_gpu = torch.cuda.is_available()
    if train_on_gpu:
        device = torch.device("cuda:0")
        pin_memory = True
        enable_amp = True

    imsize = -1
    PLOT = True
    sigma = 25
    sigma_ = sigma / 255.

    save_folder = './result'
    # deJPEG
    # fname = 'data/denoising/snail.jpg'

    ## denoising
    fname = 'data/denoising/F16_GT.png'

    # Setup
    INPUT = 'noise'  # 'meshgrid'
    pad = 'reflection'
    OPT_OVER = 'net'  # 'net,input'

    reg_noise_std = 1. / 30.  # set to 1./20. for sigma=50
    LR = 0.01

    OPTIMIZER = 'adam'  # 'LBFGS'
    show_every = 100
    exp_weight = 0.99


def load_image(config):
    fname = config.fname
    imsize = config.imsize
    PLOT = config.PLOT
    sigma_ = config.sigma_
    save_folder = config.save_folder
    # Load image
    if fname == 'data/denoising/snail.jpg':
        img_noisy_pil = crop_image(get_image(fname, imsize)[0], d=32)
        img_noisy_np = pil_to_np(img_noisy_pil)

        # As we don't have ground truth
        img_pil = img_noisy_pil
        img_np = img_noisy_np

        if PLOT:
            plot_image_grid([img_np], 4, 5)

    elif fname == 'data/denoising/F16_GT.png':
        # Add synthetic noise
        img_pil = crop_image(get_image(fname, imsize)[0], d=32)
        img_np = pil_to_np(img_pil)

        img_noisy_pil, img_noisy_np = get_noisy_image(img_np, sigma_)

        if PLOT:
            grid = plot_image_grid([img_np, img_noisy_np], 2, (8, 4))
            torchvision.utils.save_image(grid, os.path.join(save_folder, 'input.png'))
    else:
        assert False
    return img_noisy_pil, img_noisy_np, img_pil, img_np


def setup_network(config):
    fname = config.fname
    pad = config.pad
    device = config.device
    if fname == 'data/denoising/snail.jpg':
        num_iter = 2400
        input_depth = 3
        figsize = 5

        net = skip(
            input_depth, 3,
            num_channels_down=[8, 16, 32, 64, 128],
            num_channels_up=[8, 16, 32, 64, 128],
            num_channels_skip=[0, 0, 0, 4, 4],
            upsample_mode='bilinear',
            need_sigmoid=True, need_bias=True, pad=pad, act_fun='LeakyReLU')

        net = net.to(device)

    elif fname == 'data/denoising/F16_GT.png':
        num_iter = 3000
        input_depth = 32
        figsize = 4

        net = get_net(input_depth, 'skip', pad,
                      skip_n33d=128,
                      skip_n33u=128,
                      skip_n11=4,
                      num_scales=5,
                      upsample_mode='bilinear').to(device)

    else:
        assert False
    config.num_iter = num_iter
    config.input_depth = input_depth
    config.figsize = figsize
    return net


def main():
    cfg = Configuration()
    save_folder = cfg.save_folder
    os.makedirs(save_folder, exist_ok=True)

    img_noisy_pil, img_noisy_np, img_pil, img_np = load_image(cfg)
    net = setup_network(cfg)

    # input
    input_depth = cfg.input_depth
    INPUT = cfg.INPUT
    device = cfg.device
    net_input = get_noise(input_depth, INPUT, (img_pil.size[1], img_pil.size[0])).to(device).detach()

    # Compute number of parameters
    s = sum([np.prod(list(p.size())) for p in net.parameters()])
    print('Number of params: %d' % s)

    # Loss
    mse = torch.nn.MSELoss().to(device)

    img_noisy_torch = np_to_torch(img_noisy_np).to(device)

    # Optimize
    net_input_saved = net_input.detach().clone()
    noise = net_input.detach().clone()
    out_avg = None
    last_net = None
    psrn_noisy_last = 0

    # training
    OPT_OVER = cfg.OPT_OVER
    optimizer_type = cfg.OPTIMIZER
    reg_noise_std = cfg.reg_noise_std
    exp_weight = cfg.exp_weight
    show_every = cfg.show_every
    PLOT = cfg.PLOT
    figsize = cfg.figsize
    num_iter = cfg.num_iter
    LR = cfg.LR

    parameters = get_params(OPT_OVER, net, net_input)

    if optimizer_type != 'adam':
        raise NotImplementedError(optimizer_type)

    max_snr = 0

    print('Starting optimization with ADAM')
    optimizer = torch.optim.Adam(parameters, lr=LR)
    output = np.empty((num_iter, 4), np.float32)
    for i in range(num_iter):
        optimizer.zero_grad()
        if reg_noise_std > 0:
            net_input = net_input_saved + (noise.normal_() * reg_noise_std)

        out = net(net_input)

        # Smoothing
        if out_avg is None:
            out_avg = out.detach()
        else:
            out_avg = out_avg * exp_weight + out.detach() * (1 - exp_weight)

        total_loss = mse(out, img_noisy_torch)
        total_loss.backward()

        psrn_noisy = compare_psnr(img_noisy_np, out.detach().cpu().numpy()[0])
        psrn_gt = compare_psnr(img_np, out.detach().cpu().numpy()[0])
        psrn_gt_sm = compare_psnr(img_np, out_avg.detach().cpu().numpy()[0])

        # Note that we do not have GT for the "snail" example
        # So 'PSRN_gt', 'PSNR_gt_sm' make no sense
        print('Iteration %05d    Loss %f   PSNR_noisy: %f   PSRN_gt: %f PSNR_gt_sm: %f' % (
            i, total_loss.item(), psrn_noisy, psrn_gt, psrn_gt_sm), '\r', end='')
        output[i, 0] = total_loss.item()
        output[i, 1] = psrn_noisy
        output[i, 2] = psrn_gt
        output[i, 3] = psrn_gt_sm
        if PLOT and i % show_every == 0:
            out_np = torch_to_np(out)
            grid = plot_image_grid([np.clip(out_np, 0, 1),
                             np.clip(torch_to_np(out_avg), 0, 1)])

            cur_snr = max(psrn_gt, psrn_gt_sm)
            if max_snr < cur_snr:
                max_snr = cur_snr
                torchvision.utils.save_image(grid, os.path.join(save_folder, 'betterPNSR.png'))

                out_np.tofile(os.path.join(save_folder, 'better_gt.dat'))
                out_avg.detach().cpu().numpy().tofile(os.path.join(save_folder, 'better_gt_sm.dat'))


        # Backtracking
        if i % show_every:
            if psrn_noisy - psrn_noisy_last < -5:
                print('Falling back to previous checkpoint.')

                for new_param, net_param in zip(last_net, net.parameters()):
                    net_param.data.copy_(new_param.cuda())

                return total_loss * 0
            else:
                last_net = [x.detach().cpu() for x in net.parameters()]
                psrn_noisy_last = psrn_noisy

        optimizer.step()

    out_np = torch_to_np(net(net_input))
    grid = plot_image_grid([np.clip(out_np, 0, 1), img_np])
    torchvision.utils.save_image(grid, os.path.join(save_folder, 'FinalOutput.png'))
    out_np.tofile(os.path.join(save_folder, 'final_out.dat'))


    labels = ['loss', 'PNSR noisy', 'PSNR gt', 'PNSR gt sm']
    np.savetxt(os.path.join(save_folder, 'loss.txt'), output, '%f', header=', '.join(labels))
    view_metric(os.path.join(save_folder, 'loss.txt'))
    pass


def view_metric(file):
    output = np.loadtxt(file, np.float32)
    save_folder = os.path.dirname(file)

    labels = ['loss', 'PNSR noisy', 'PSNR gt', 'PNSR gt sm']

    fig = plt.figure(figsize=(8, 4))

    ax1 = fig.add_subplot(121)
    ax2 = fig.add_subplot(122)

    for i, label in enumerate(labels):
        if i == 0:
            ax1.plot(output[:, i], color=color_list[i], label=labels[i], linewidth=linewidth)
        else:
            ax2.plot(output[:, i], color=color_list[i], label=labels[i], linewidth=linewidth)


    ax1.set_xlabel('Epoch #', fontdict=font)
    ax1.set_ylabel("MSE", fontdict=font)
    ax1.legend(prop=font)

    temp = ax1.get_xticklabels()
    ax1.set_xticklabels(temp, **font)
    temp = ax1.get_yticklabels()
    ax1.set_yticklabels(temp, **font)

    ax2.set_xlabel('Epoch #', fontdict=font)
    ax2.set_ylabel('PSNR', fontdict=font)
    ax2.legend(prop=font)

    temp = ax2.get_xticklabels()
    ax2.set_xticklabels(temp, **font)
    temp = ax2.get_yticklabels()
    ax2.set_yticklabels(temp, **font)
    # plt.xticks(**font)
    # plt.yticks(**font)

    plt.suptitle(f'Deep image prior', fontdict=font)
    plt.savefig(os.path.join(save_folder, 'loss.png'), bbox_inches='tight')



if __name__ == '__main__':
    # view_metric('./result/loss.txt')
    main()
    print('done')
